# 导入必要的库
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

# 检查是否有可用的GPU，如果有则使用GPU加速
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# ----------------------- 超参数设置 -----------------------
batch_size = 64        # 每个mini-batch的样本数
learning_rate = 0.001  # 学习率
num_epochs = 10        # 训练轮数

# ----------------------- 数据预处理 -----------------------
# 定义数据预处理操作，将图像转换为Tensor并进行归一化处理
# MNIST数据集的均值为0.1307，标准差为0.3081
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 下载MNIST训练集和测试集
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)

# 使用DataLoader加载数据集，方便后续mini-batch训练
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)

# ----------------------- 定义卷积神经网络 -----------------------
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 第一层卷积：输入通道数为1（灰度图），输出通道数为32，卷积核大小为3，步长为1，边缘补零1
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=1, padding=1)
        # 第二层卷积：输入32，输出64，卷积核大小为3，步长为1，边缘补零1
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1)
        # 池化层：使用2x2的最大池化，降低特征图的尺寸
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        # Dropout层：防止过拟合，随机将部分神经元置0
        self.dropout = nn.Dropout(0.25)
        # 全连接层1：将卷积层提取的特征展平后输入，输入特征数为64 * 14 * 14，输出特征数为128
        self.fc1 = nn.Linear(in_features=64 * 14 * 14, out_features=128)
        # 全连接层2：输出层，将128维特征映射到10个类别（数字0-9）
        self.fc2 = nn.Linear(in_features=128, out_features=10)

    def forward(self, x):
        # 第1层卷积 + 激活函数ReLU
        x = torch.relu(self.conv1(x))
        # 第2层卷积 + 激活函数ReLU
        x = torch.relu(self.conv2(x))
        # 池化层，降低特征图尺寸
        x = self.pool(x)
        # 使用Dropout防止过拟合
        x = self.dropout(x)
        # 将多维张量展平为二维张量，方便输入全连接层
        x = x.view(x.size(0), -1)
        # 全连接层1 + ReLU激活
        x = torch.relu(self.fc1(x))
        # 输出层（不需要激活函数，因为交叉熵损失函数内部包含了softmax操作）
        x = self.fc2(x)
        return x

# 实例化模型，并将模型移动到设备上（GPU或CPU）
model = CNN().to(device)

# ----------------------- 定义损失函数和优化器 -----------------------
# 交叉熵损失函数适用于多分类问题
criterion = nn.CrossEntropyLoss()
# 使用Adam优化器更新网络参数
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# ----------------------- 模型训练 -----------------------
print("开始训练...")
for epoch in range(num_epochs):
    model.train()  # 将模型设置为训练模式
    running_loss = 0.0  # 用于统计每100个mini-batch的损失和
    for i, (images, labels) in enumerate(train_loader):
        # 将数据移动到设备上
        images, labels = images.to(device), labels.to(device)

        # 前向传播：将输入数据传入模型，得到输出
        outputs = model(images)
        # 计算预测结果与真实标签之间的损失
        loss = criterion(outputs, labels)

        # 反向传播：先将梯度置零，再进行梯度计算和参数更新
        optimizer.zero_grad()  # 清空上一步的梯度信息
        loss.backward()        # 反向传播计算梯度
        optimizer.step()       # 更新模型参数

        running_loss += loss.item()  # 累计损失值
        # 每训练100个mini-batch，打印一次平均损失
        if (i + 1) % 100 == 0:
            print(f'轮数 [{epoch+1}/{num_epochs}], 第 [{i+1}/{len(train_loader)}] 个mini-batch, 平均损失: {running_loss/100:.4f}')
            running_loss = 0.0

# ----------------------- 模型测试 -----------------------
model.eval()  # 将模型设置为评估模式，关闭Dropout等训练专用模块
correct = 0  # 统计预测正确的样本数量
total = 0    # 统计测试集中样本总数

# 测试过程中不需要计算梯度，节省内存和计算资源
with torch.no_grad():
    for images, labels in test_loader:
        # 将数据移动到设备上
        images, labels = images.to(device), labels.to(device)
        # 计算模型输出
        outputs = model(images)
        # 取得预测结果（每个样本的输出向量中取最大值对应的类别）
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)  # 更新总样本数
        correct += (predicted == labels).sum().item()  # 累计预测正确的样本数

# 输出测试集上的准确率
print(f'测试集上的准确率为: {100 * correct / total:.2f}%')
